# encoding: utf-8


"""

@author: tongzhenguo

@time: 2021/6/7 上午11:36

@desc:

封装mongo操作
"""

import collections
import traceback

import pandas as pd
import pymongo


class MongoDB:
    def __init__(self, **dbconfig):
        self.dbconfig = dbconfig
        try:
            # 默认最大连接池100个
            self.mg_client = pymongo.MongoClient(self.dbconfig['host'], self.dbconfig['port'], maxPoolSize=5)
            self.mg_client.admin.authenticate(self.dbconfig['user'], self.dbconfig['password'])
            self.db = self.mg_client[self.dbconfig['database']]
            print("mongo authenticate sucess ====")
        except Exception as e:
            print("mongo connect error , msg : %s" % repr(e))
        finally:
            if self.mg_client: self.mg_client.close()

    def fetch_data(self, column_list, collection_name, query_dict):
        """从mongo中提取数据"""
        try:
            # 组装成一个key是colomn,列是值列表的大dict
            column_map = collections.defaultdict(list)
            for item in self.db[collection_name].find(query_dict):
                record = []
                for col in column_list:
                    record.append(None if col not in item else item[col])
                for idx, col in enumerate(column_list):
                    column_map[col].append(record[idx])
                if len(column_map[column_list[0]]) % 100 == 0:
                    print("process data count : %s " % (len(column_map[column_list[0]])))
            return pd.DataFrame(column_map)
        except Exception as e:
            print("mongo fetch_data error , msg : %s" % repr(e))
            traceback.print_exc()

    def exists_one(self, collection_name, query_dict):
        one = self.db[collection_name].find_one(query_dict)
        return False if not one else True

    def add(self, table_name, document):
        """向table插入数据"""
        assert type(document) == dict
        return self.db[table_name].insert_one(document).inserted_id

    def upsert(self, table_name, where, document):
        return self.db[table_name].update_one(where, {"$set": document}, upsert=True).modified_count

    def delete(self, table_name, where):
        """删除操作"""
        return self.db[table_name].delete_many(where).deleted_count


if __name__ == "__main__":
    mongo_db = MongoDB(uri="mongodb://127.0.0.1:27017", database="stock")
    # TODO 添加测试case
    # mongo_db.add("stock", {"_id": "700001", "close": 70})
    # 会报错E11000 duplicate key error collection
    # mongo_db.add("stock", {"_id": "000001", "close": 71})
    mongo_db.upsert("stock", {"_id": "700001"}, {"close2": 71})
    # 任务状态，为1代表已存在
    # mongo_db.add("collect_task_info", {"_id": "000001-日频数据_估值指标-20210609", "done": 1})
    # mongo_db.delete("stock", {"_id": "000001"})
    # mongo_db.delete("stock", {})
    # mongo_db.delete("collect_task_info", {"_id": "000001-日频数据_估值指标-20210609", "done": 1})
